import sys
sys.path.append('./pycocotools')
from pycocotools.coco import COCO
import numpy as np
import json


def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = box[0] + (box[2] / 2.0)
    y = box[1] + (box[3] / 2.0)
    w = box[2]
    h = box[3]
    x = x * dw
    y = y * dh
    w = w * dw
    h = h * dh
    return x, y, w, h


def save_json(output_dir, json_file):
    with open(output_dir, 'w') as json_f:
        json_f.write('[\n')
        n = 0
        for each_dict in json_file:
            n = n + 1
            json_f.write('    ' + json.dumps(each_dict, ensure_ascii=False))
            if n != len(json_file):
                json_f.write(',')
            json_f.write('\n')
        json_f.write(']')


if __name__ == "__main__":
    dataDir = '/dataset/COCO/2017'
    dataType = 'train2017'
    annFile = '{}/annotations/instances_{}.json'.format(dataDir, dataType)

    coco = COCO(annFile)

    coco_label = coco.cats

    i = 0
    label_map = {}
    coco_json =[]
    for k, v in coco_label.items():
        label_map[v["id"]] = i
        v["id"] = i
        coco_json.append(v)
        i += 1

    json_out = '{}/coco_list/coco_class.json'.format(dataDir)
    save_json(json_out, coco_json)

    img_path_list = []
    img_bbox_list = []
    img_size_list = []

    for key, value in coco.imgs.items():
        annIds = coco.getAnnIds(imgIds=key)
        img_path = '{}/images/{}/{}'.format(dataDir, dataType, value['file_name'])
        img_h = value['height']
        img_w = value['width']

        bboxes = []
        for ann_id in annIds:
            bbox = []
            label = label_map[coco.anns[ann_id]['category_id']]
            bbox.append(label)
            if label > 80:
                print('oh no')
            tmp = coco.anns[ann_id]['bbox']
            tmp = convert((img_w, img_h), tmp)
            bbox.append(tmp[0])
            bbox.append(tmp[1])
            bbox.append(tmp[2])
            bbox.append(tmp[3])
            bboxes.append(bbox)

        if len(bboxes) > 0:
            img_path_list.append(img_path)
            img_bbox_list.append(bboxes)
            img_size_list.append((img_h, img_w))

    datas = np.array([
        np.array([
            np.array(img_path_list[i]),
            np.array(img_bbox_list[i]),
            np.array(img_size_list[i])]
        ) for i in range(len(img_path_list))])

    output_file = './data/coco/coco_img_ann.npy'
    np.save(output_file, datas)

